// Multi-Agent 系统核心实现
import { ChatOllama } from "@langchain/ollama";
import { ChatPromptTemplate } from "@langchain/core/prompts";
import { StringOutputParser } from "@langchain/core/output_parsers";
import { RunnableSequence } from "@langchain/core/runnables";
import { ToolManager } from "./tool-manager";

// 1. 智能体角色信息（类似“岗位说明书”）
interface AgentRole {
  name: string; // 角色名称（如“研究员”）
  description: string; // 角色职责
  expertise: string[]; // 专长（如“搜索”“数据分析”）
  tools: string[]; // 可使用的工具名称（如“search”“calculator”）
}

// 2. 智能体单轮输出（类似“工作汇报”格式）
interface AgentOutput {
  role: string; // 输出者角色
  round: number; // 当前轮次
  reasoning: string; // 思考过程（“我为什么这么做”）
  suggestion: string; // 具体建议或结论
  toolCalls?: ToolCall[]; // 调用的工具及结果（可选）
  status: "continue" | "final"; // 状态：继续协作/给出最终答案
  handoff?: string | null; // 任务移交目标（如“分析师”，可选）
  finalAnswer?: string; // 最终答案（仅status为final时有效）
  confidence?: number; // 信心度（0-1，可选）
}

// 3. 工具调用记录（类似“工具使用日志”）
interface ToolCall {
  toolName: string; // 工具名（如“search”）
  args: any; // 工具参数（如搜索关键词）
  result: string; // 工具返回结果
}

// 4. 系统最终结果（类似“项目总结报告”）
interface MultiAgentResult {
  finalAnswer: string; // 最终答案
  from: string; // 给出答案的智能体
  confidence?: number; // 信心度
  rounds: number; // 总轮次
  transcript: AgentOutput[]; // 完整交互记录
}

// 5. 系统配置（类似“项目参数设置”）
interface MultiAgentOptions {
  baseUrl?: string; // 模型接口地址（如Ollama的地址）
  model?: string; // 所用模型（如“llama3.1:latest”）
  temperature?: number; // 模型随机性（0-1，越高越灵活）
  maxRounds?: number; // 最大协作轮次（防止无限循环）
  contextWindow?: number; // 共享上下文窗口大小（保留最近多少轮记录）
}

/**
 * 预定义的 Agent 角色
 */
export const AGENT_ROLES: AgentRole[] = [
  {
    name: "研究员",
    description: "负责信息收集和事实核查",
    expertise: ["搜索", "验证", "总结"],
    tools: ["search", "get_time"]
  },
  {
    name: "分析师",
    description: "负责数据分析和洞察",
    expertise: ["计算", "统计", "可视化"],
    tools: ["calculator", "analyze_data"]
  },
  {
    name: "协调员",
    description: "负责任务分配和结果整合",
    expertise: ["规划", "协调", "总结"],
    tools: ["get_time"]
  }
];

/**
 * 默认配置
 */
const DEFAULT_OPTIONS: Required<MultiAgentOptions> = {
  baseUrl: "http://192.168.31.159:11434",
  model: "llama3.1:latest",
  temperature: 0.7,
  maxRounds: 5,
  contextWindow: 6,
};

/**
 * 单个协作 Agent
 */
class CollaborativeAgent {
  private readonly role: AgentRole;
  private readonly chain: RunnableSequence<Record<string, unknown>, string>;
  private readonly toolManager: ToolManager;

  constructor(role: AgentRole, toolManager: ToolManager, config: Required<MultiAgentOptions>) {
    this.role = role;
    this.toolManager = toolManager;
    this.chain = this.initializeChain(role, config);
  }

  /**
   * 生成响应
   */
  async generateResponse(input: {
    task: string;
    round: number;
    phaseHint: string;
    sharedContext: string;
    personalHistory: string;
  }): Promise<AgentOutput> {
    // 调用模型获取原始输出
    const raw = await this.chain.invoke({
      task: input.task,
      round: input.round,
      phase_hint: input.phaseHint,
      shared_context: input.sharedContext,
      personal_history: input.personalHistory,
    });

    // 解析输出为结构化数据
    const parsed = this.parseOutput(raw);

    // 如果需要调用工具
    const toolCalls: ToolCall[] = [];
    if (parsed.toolCalls && parsed.toolCalls.length > 0) {
      for (const call of parsed.toolCalls) {
        try {
          const result = await this.toolManager.callTool(call.toolName, call.args);
          toolCalls.push({ ...call, result });
        } catch (error) {
          toolCalls.push({
            ...call,
            result: `工具调用失败: ${error instanceof Error ? error.message : '未知错误'}`
          });
        }
      }
    }

    return {
      role: this.role.name,
      round: input.round,
      reasoning: parsed.reasoning,
      suggestion: parsed.suggestion,
      toolCalls: toolCalls.length > 0 ? toolCalls : undefined,
      status: parsed.status,
      handoff: parsed.handoff,
      finalAnswer: parsed.finalAnswer,
      confidence: parsed.confidence,
    };
  }

  /**
   * 初始化 Agent 链
   */
  private initializeChain(role: AgentRole, config: Required<MultiAgentOptions>) {
    const systemPrompt = `你是 ${role.name}，${role.description}。

你的专长：${role.expertise.join("、")}
可用工具：${role.tools.length ? role.tools.join("、") : "暂无"}

在每一轮协作中，你需要：
1. 阅读共享上下文和你的历史记录
2. 基于你的专长提供分析和建议
3. 如果需要，调用工具获取信息
4. 决定是继续协作还是给出最终答案

输出格式要求（严格遵守）：
思考：<你的分析过程>
建议：<你的建议或结论>
工具调用：<如果需要调用工具，格式为 [工具名](参数JSON)，例如 [search]({"query":"人工智能"})，可以多个>
状态：<continue 或 final>
移交：<如果需要其他 Agent 接手，写出角色名，否则写 null>
最终答案：<仅当状态为 final 时填写>
信心度：<0-1 之间的数字，可选>

注意：
- 如果你无法独立完成任务，应该移交给更合适的 Agent
- 如果所有必要信息都已收集，协调员应该给出最终答案
- 工具调用格式示例：[calculator]({"expression":"10+20"})
`;

    const humanPrompt = `任务：{task}
当前轮次：第 {round} 轮
阶段说明：{phase_hint}

共享上下文：
{shared_context}

你的历史记录：
{personal_history}

请基于以上信息给出你的分析和建议：`;

    const prompt = ChatPromptTemplate.fromMessages([
      ["system", systemPrompt.trim()],
      ["human", humanPrompt.trim()],
    ]);

    const model = new ChatOllama({
      baseUrl: config.baseUrl,
      model: config.model,
      temperature: config.temperature,
    });

    return RunnableSequence.from([
      prompt,
      model,
      new StringOutputParser(),
    ]);
  }

  /**
   * 解析 Agent 输出
   */
  private parseOutput(raw: string): {
    reasoning: string;
    suggestion: string;
    toolCalls?: ToolCall[];
    status: "continue" | "final";
    handoff?: string | null;
    finalAnswer?: string;
    confidence?: number;
  } {
    const result: any = {
      reasoning: "",
      suggestion: "",
      status: "continue",
      handoff: null,
    };

    // 解析思考
    const thinkingMatch = raw.match(/思考\s*[：:]\s*(.*?)(?=\n(?:建议|工具调用|状态|移交|最终答案)|$)/s);
    if (thinkingMatch) {
      result.reasoning = thinkingMatch[1].trim();
    }

    // 解析建议
    const suggestionMatch = raw.match(/建议\s*[：:]\s*(.*?)(?=\n(?:工具调用|状态|移交|最终答案)|$)/s);
    if (suggestionMatch) {
      result.suggestion = suggestionMatch[1].trim();
    }

    // 解析工具调用 [toolName](args)
    const toolCallPattern = /\[(\w+)\]\(({[^)]+})\)/g;
    const toolCalls: ToolCall[] = [];
    let match;
    while ((match = toolCallPattern.exec(raw)) !== null) {
      try {
        const toolName = match[1];
        const args = JSON.parse(match[2]);
        toolCalls.push({ toolName, args, result: "" });
      } catch (e) {
        console.warn("工具调用解析失败:", match[0]);
      }
    }
    if (toolCalls.length > 0) {
      result.toolCalls = toolCalls;
    }

    // 解析状态
    const statusMatch = raw.match(/状态\s*[：:]\s*(continue|final)/i);
    if (statusMatch) {
      result.status = statusMatch[1].toLowerCase();
    }

    // 解析移交
    const handoffMatch = raw.match(/移交\s*[：:]\s*(.*?)(?=\n|$)/);
    if (handoffMatch) {
      const handoff = handoffMatch[1].trim();
      result.handoff = (handoff === "null" || handoff === "无") ? null : handoff;
    }

    // 解析最终答案
    const finalAnswerMatch = raw.match(/最终答案\s*[：:]\s*(.*?)(?=\n(?:信心度)|$)/s);
    if (finalAnswerMatch) {
      result.finalAnswer = finalAnswerMatch[1].trim();
    }

    // 解析信心度
    const confidenceMatch = raw.match(/信心度\s*[：:]\s*(0?\.\d+|1\.0|1)/);
    if (confidenceMatch) {
      result.confidence = parseFloat(confidenceMatch[1]);
    }

    return result;
  }
}

/**
 * Multi-Agent 协调器
 */
export class MultiAgentOrchestrator {
  private readonly config: Required<MultiAgentOptions>;
  private readonly toolManager: ToolManager;
  private readonly agents: Map<string, CollaborativeAgent> = new Map();
  private readonly executionOrder: string[] = [];
  private transcript: AgentOutput[] = [];
  private personalHistories: Map<string, AgentOutput[]> = new Map();

  constructor(options: MultiAgentOptions = {}) {
    this.config = {
      ...DEFAULT_OPTIONS,
      ...options,
    };

    this.toolManager = new ToolManager();
    this.initializeAgents();
  }

  /**
   * 运行 Multi-Agent 系统
   */
  async run(task: string): Promise<MultiAgentResult> {
    console.log(`\n🚀 Multi-Agent 系统启动`);
    console.log(`📋 任务: ${task}\n`);

    this.transcript = [];
    this.personalHistories = new Map(
      AGENT_ROLES.map((role) => [role.name, [] as AgentOutput[]]),
    );

    let finalResult: MultiAgentResult | null = null;

    for (let round = 1; round <= this.config.maxRounds; round++) {
      console.log(`\n━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━`);
      console.log(`🔄 第 ${round} 轮协作`);
      console.log(`━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\n`);

      const visitedThisRound = new Set<string>();
      const queue = [...this.executionOrder];

      while (queue.length) {
        const roleName = queue.shift()!;
        const agent = this.agents.get(roleName);
        if (!agent) {
          continue;
        }

        if (visitedThisRound.has(roleName)) {
          continue;
        }
        visitedThisRound.add(roleName);

        console.log(`\n👤 ${roleName} 开始工作...`);

        const turn = await agent.generateResponse({
          task,
          round,
          phaseHint: this.buildPhaseHint(roleName, round),
          sharedContext: this.buildSharedContext(roleName),
          personalHistory: this.buildPersonalHistory(roleName),
        });

        this.recordTurn(turn);
        this.printTurn(turn);

        // 检查是否给出最终答案
        if (turn.status === "final" && turn.finalAnswer) {
          finalResult = {
            finalAnswer: turn.finalAnswer,
            from: roleName,
            confidence: turn.confidence,
            rounds: round,
            transcript: [...this.transcript],
          };
          console.log(`\n✅ ${roleName} 给出了最终答案！`);
          break;
        }

        // 处理移交
        if (turn.handoff) {
          const target = turn.handoff.trim();
          if (target && this.agents.has(target) && !visitedThisRound.has(target)) {
            console.log(`\n🔀 ${roleName} 将任务移交给 ${target}`);
            queue.unshift(target);
          }
        }
      }

      if (finalResult) {
        break;
      }
    }

    // 如果没有得到最终答案，由协调员强制总结
    if (!finalResult) {
      console.log(`\n⚠️ 达到最大轮次，由协调员进行最终总结...`);
      const coordinator = this.agents.get("协调员");
      if (coordinator) {
        const finalTurn = await coordinator.generateResponse({
          task,
          round: this.config.maxRounds + 1,
          phaseHint: "请基于所有讨论给出最终总结",
          sharedContext: this.buildSharedContext("协调员"),
          personalHistory: "",
        });

        finalResult = {
          finalAnswer: finalTurn.finalAnswer || finalTurn.suggestion || "未能得出最终答案",
          from: "协调员",
          confidence: finalTurn.confidence,
          rounds: this.config.maxRounds,
          transcript: [...this.transcript, finalTurn],
        };
      } else {
        finalResult = {
          finalAnswer: "未能得出最终答案",
          from: "系统",
          rounds: this.config.maxRounds,
          transcript: [...this.transcript],
        };
      }
    }

    console.log(`\n━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━`);
    console.log(`🎯 最终结果`);
    console.log(`━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━`);
    console.log(`\n来自: ${finalResult.from}`);
    console.log(`轮次: ${finalResult.rounds}`);
    if (finalResult.confidence) {
      console.log(`信心度: ${(finalResult.confidence * 100).toFixed(0)}%`);
    }
    console.log(`\n答案:\n${finalResult.finalAnswer}\n`);

    return finalResult;
  }

  /**
   * 初始化所有 Agent
   */
  private initializeAgents() {
    for (const role of AGENT_ROLES) {
      const agent = new CollaborativeAgent(role, this.toolManager, this.config);
      this.agents.set(role.name, agent);
      this.executionOrder.push(role.name);
    }
    console.log("✅ 已初始化 Agents:", this.executionOrder.join(", "));
  }

  /**
   * 记录 Agent 输出
   */
  private recordTurn(turn: AgentOutput) {
    this.transcript.push(turn);
    const history = this.personalHistories.get(turn.role);
    if (history) {
      history.push(turn);
    } else {
      this.personalHistories.set(turn.role, [turn]);
    }
  }

  /**
   * 打印 Agent 输出
   */
  private printTurn(turn: AgentOutput) {
    console.log(`\n💭 思考: ${turn.reasoning}`);
    console.log(`💡 建议: ${turn.suggestion}`);
    
    if (turn.toolCalls && turn.toolCalls.length > 0) {
      console.log(`\n🔧 工具调用:`);
      turn.toolCalls.forEach(call => {
        console.log(`  - ${call.toolName}(${JSON.stringify(call.args)})`);
        console.log(`    结果: ${call.result}`);
      });
    }

    console.log(`\n📊 状态: ${turn.status}`);
    if (turn.handoff) {
      console.log(`🔀 移交: ${turn.handoff}`);
    }
    if (turn.finalAnswer) {
      console.log(`\n✨ 最终答案: ${turn.finalAnswer}`);
    }
    if (turn.confidence) {
      console.log(`📈 信心度: ${(turn.confidence * 100).toFixed(0)}%`);
    }
  }

  /**
   * 构建阶段提示
   */
  private buildPhaseHint(roleName: string, round: number): string {
    const hints: Record<string, string> = {
      "研究员": "请优先收集事实、列出关键证据，并指出信息缺口。",
      "分析师": "请基于当前事实进行推理、计算或验证，并提出洞察。",
      "协调员": "请统筹各方观点，指出冲突与共识，推动形成结论。",
    };

    const hint = hints[roleName] || "请基于你的专长提供建议。";
    
    if (round === 1) {
      return `初始分析阶段。${hint}`;
    } else if (round >= this.config.maxRounds - 1) {
      return `接近尾声，请尽快推动形成最终结论。${hint}`;
    } else {
      return `协作调整阶段。${hint}`;
    }
  }

  /**
   * 构建共享上下文
   */
  private buildSharedContext(currentRole: string): string {
    const recentTurns = this.transcript.slice(-this.config.contextWindow);
    
    if (recentTurns.length === 0) {
      return "暂无历史讨论";
    }

    const context = recentTurns
      .filter(turn => turn.role !== currentRole)
      .map(turn => {
        let text = `[${turn.role}] ${turn.suggestion}`;
        if (turn.toolCalls && turn.toolCalls.length > 0) {
          text += `\n  工具结果: ${turn.toolCalls.map(c => c.result).join("; ")}`;
        }
        return text;
      })
      .join("\n\n");

    return context || "暂无其他 Agent 的意见";
  }

  /**
   * 构建个人历史
   */
  private buildPersonalHistory(roleName: string): string {
    const history = this.personalHistories.get(roleName) || [];
    
    if (history.length === 0) {
      return "这是你的第一次参与";
    }

    const recentHistory = history.slice(-3);
    return recentHistory
      .map(turn => `[第${turn.round}轮] ${turn.reasoning} → ${turn.suggestion}`)
      .join("\n");
  }
}
